;; Matrix chain multiplication
;; Copyright (C) 2022  Rocks Mazama
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

(ns mazama.linalg.matrix-chain)

(defn- cost
  [fnrows fncols A1 A2]
  (* (fnrows A1) (fnrows A2) (fncols A2)))

(def ^:private helper nil)

(defn- algorithm
  [fnrows fncols]
  (def ^:private helper
    (memoize (fn [coll]
               (case (count coll)
                 0 {::chain coll ::cost 0 ::ncols 0 ::nrows 0}
                 1 (let [A (first coll)] {::chain A ::cost 0 ::ncols (fncols A) ::nrows (fnrows A)})
                 2 {::chain coll
                    ::cost (cost fnrows fncols (first coll) (second coll))
                    ::ncols (fncols (second coll))
                    ::nrows (fnrows (first coll))}
                 (loop [result {::cost ##Inf}
                        splits (map #(split-at % coll) (range 1 (count coll)))]
                   (if (empty? splits)
                     result
                     (let [[left-split right-split] (first splits)
                           left-chain (helper left-split)]
                       (if (>= (::cost left-chain) (::cost result))
                         (recur result (rest splits))
                         (let [right-chain (helper right-split)]
                           (if (>= (::cost right-chain) (::cost result))
                             (recur result (rest splits))
                             (let [cost-chain (+ (::cost left-chain)
                                                 (::cost right-chain)
                                                 (cost ::nrows ::ncols left-chain right-chain))]
                               (if (>= cost-chain (::cost result))
                                 (recur result (rest splits))
                                 (recur {::chain (vector (::chain left-chain) (::chain right-chain))
                                         ::cost cost-chain
                                         ::ncols (::ncols right-chain)
                                         ::nrows (::nrows left-chain)}
                                        (rest splits))))))))))))))
  helper)

(defn matrix-chain
  "A matrix chain is a parenthesized version of a sequence of matrices.
  A matrix chain is either a single matrix or a sequence of two elements,
  each element being a matrix chain.
  `fnrows` is a function that returns the number of rows in a matrix.
  `fncols` is a function that returns the number of columns in a matrix.
  `coll` is a sequence of matrices.
  It is assumed that adjacent elements of `coll` are compatible for matrix multiplication."
  [fnrows fncols coll]
  (let [result ((algorithm fnrows fncols) coll)]
    (def ^:private helper nil)
    (::chain result)))

(defn multiply
  "Multiply a matrix chain.
  `f*` is a function that multiplies two matrices.
  `chain` is a matrix chain."
  [f* chain]
  (if (sequential? chain)
    (let [[left right] chain]
      (f* (multiply f* left) (multiply f* right)))
    chain))
